# @日期：2021/1/17 11:16
# @公司：阿尔法工场
# @作者：张大鹏
# @昵称：lxgzhw
# @文件：01.简单线性回归根据受教育年限预测收入.py
# @描述：简单线性回归根据受教育年限预测收入

# 引入依赖
import torch
# pytorch的核心高级API
from torch import nn
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 读取数据集
df = pd.read_csv('income.csv')
# print(df.head())

# 绘图，查看受教育年限和收入之间的关系
plt.scatter(df['Education'], df['Income'])
plt.xlabel('Education')
plt.ylabel('Income')
plt.title('Education and Income')
# plt.show()

# 提取数据
X = df['Education'].values  # 使用values提取里面的值
# print(X)
# 将数据转换为1列，使用-1自动计算有多少行
X = X.reshape(-1, 1)
# 将数据类型转换为np.float32类型
X = X.astype(np.float32)
# 将numpy的array类型转换为torch的tensor张量类型
X = torch.from_numpy(X)
# print(X)

# 同样的方法，得到Y的张量
Y = df['Income'].values
Y = torch.from_numpy(Y.reshape(-1, 1).astype(np.float32))
# print(Y)

# 创建线性模型
# 参数1：输入的特征长度（一个受教育年限数据）
# 参数2：输出的特征长度（一个收入预测数据）
model = nn.Linear(1, 1)

# 计算均方误差
loss_fn = nn.MSELoss()  # 是一个类，需要加括号

# 随机梯度下降算法
# model.parameters：返回模型需要优化的参数
# lr：学习速率
opt = torch.optim.SGD(model.parameters(), lr=0.0001)

# 训练模型
for epoch in range(5000):  # 对全部数据训练5000次
    for x, y in zip(X, Y):  # 同时对X和Y进行迭代
        # 放入模型预测
        y_pred = model(x)
        # 计算损失
        # 参数1：真实值
        # 参数2：预测值
        loss = loss_fn(y, y_pred)
        # 清空原有的梯度
        opt.zero_grad()
        # 反向传播，求解梯度
        loss.backward()
        # 执行优化：优化模型参数
        opt.step()

# 查看训练好的模型的权重
# print(model.weight)
# 查看训练好的模型的偏置
# print(model.bias)

# 绘制图片查看
# X.numpy()：将tensor张量转换numpy数组
# model(X)：关于X的预测结果
# model(X).data.numpy()：预测结果的值
# c='r'：红色
plt.plot(X.numpy(), model(X).data.numpy(), c='r')
# 保存图片：注意，如果不调用show，无法保存图片
plt.savefig('EducationAndIncome.png')
plt.show()

# 保存模型
torch.save(model.state_dict(), "EducationAndIncome.model")
# 加载模型
model.load_state_dict(torch.load("EducationAndIncome.model"))
model.eval()
